import numpy as np
import matplotlib.pyplot as plt


def classical_pendulum(theta0, b=0.05, t_max=10):
    """模拟经典阻尼单摆运动（欧拉法）"""
    dt = 0.01
    g = 9.8  # 重力加速度(m/s²)
    L = 1.0  # 摆线长度(m)

    t = np.arange(0, t_max, dt)
    theta, omega = [theta0], [0]  # 初始角度和角速度

    for i in range(1, len(t)):
        # 运动方程: d²θ/dt² + b·dθ/dt + (g/L)sinθ = 0
        alpha = - (g / L) * np.sin(theta[-1]) - b * omega[-1]  # 角加速度
        new_omega = omega[-1] + alpha * dt
        new_theta = theta[-1] + new_omega * dt
        omega.append(new_omega)
        theta.append(new_theta)

    return theta[:len(t)]  # 返回角度序列


# 量子不确定性模拟
def quantum_uncertainty(theta0, uncertainty_factor, num_paths=100):
    all_trajectories = []
    for _ in range(num_paths):
        # 添加初始角度量子涨落
        perturbed_theta0 = theta0 + uncertainty_factor * np.random.normal()
        traj = classical_pendulum(perturbed_theta0)
        all_trajectories.append(traj)

    # 绘制轨迹云团
    plt.figure(figsize=(10, 6))
    for i, traj in enumerate(all_trajectories):
        # 前20条轨迹用不同颜色突出显示
        color = 'red' if i < 20 else 'blue'
        alpha = 0.6 if i < 20 else 0.1
        plt.plot(traj, alpha=alpha, color=color, linewidth=0.8)

    # 添加理论说明
    plt.title(f"Quantum Uncertainty Simulation (σ={uncertainty_factor})")
    plt.xlabel('Time step (0.01s/step)')
    plt.ylabel('Pendulum angle (rad)')
    plt.annotate('Initial quantum fluctuations\ncause trajectory divergence',
                 xy=(50, theta0 + 0.05), xytext=(200, 0.3),
                 arrowprops=dict(facecolor='black', shrink=0.05))
    plt.grid(alpha=0.3)
    plt.show()


# 示例调用
quantum_uncertainty(theta0=0.2, uncertainty_factor=0.1)
